import os
from typing import Dict
import logging

def ensure_dir(path: str):
    os.makedirs(path, exist_ok=True)

def format_memory(mb: int) -> str:
    if mb >= 1024:
        return f"{mb/1024:.1f} GB"
    return f"{mb} MB"

def get_actual_gpu_usage(device: str) -> float:
    import pynvml
    handle = pynvml.nvmlDeviceGetHandleByIndex(int(device.split(":")[-1]))
    return pynvml.nvmlDeviceGetMemoryInfo(handle).used

def setup_logging(log_file: str):
    # 设置日志级别、格式和日志文件路径
    logging.basicConfig(
        level=logging.DEBUG,  # 设置日志级别
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',  # 日志格式
        filename=log_file,  # 日志文件路径
        filemode='a'  # 写入模式 'w'表示覆盖，'a'表示追加
    )
